{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "slXNX2PTqC-D"
   },
   "source": [
    "<center>\n",
    "    <p style=\"text-align:center\">\n",
    "        <img alt=\"phoenix logo\" src=\"https://storage.googleapis.com/arize-phoenix-assets/assets/phoenix-logo-light.svg\" width=\"200\"/>\n",
    "        <br>\n",
    "        <a href=\"https://arize.com/docs/phoenix/\">Docs</a>\n",
    "        |\n",
    "        <a href=\"https://github.com/Arize-ai/phoenix\">GitHub</a>\n",
    "        |\n",
    "        <a href=\"https://arize-ai.slack.com/join/shared_invite/zt-2w57bhem8-hq24MB6u7yE_ZF_ilOYSBw#/shared-invite/email\">Community</a>\n",
    "    </p>\n",
    "</center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lNmSeMhTp47g"
   },
   "source": [
    "# LangGraph Router Pattern: Intent-Based Agent Routing\n",
    "In this tutorial, we'll explore the Router pattern using LangGraph, a framework for building dynamic, stateful LLM applications. We’ll build a smart support assistant that can route customer queries to the appropriate specialized agent—Billing, Tech Support, or General Info—based on the user’s intent.\n",
    "\n",
    "LangGraph enables this by letting us define a structured graph of logic, where a router node classifies the input, and conditional edges forward it to the correct sub-agent. Each agent uses local context (e.g., invoice history, troubleshooting tips) to respond intelligently.\n",
    "\n",
    "We also trace this application using Phoenix, which gives us complete visibility into routing decisions, tool usage, and model interactions. This is helpful for debugging routing accuracy and understanding how the graph executes end-to-end."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install langgraph langchain langchain_community \"arize-phoenix==9.0.1\" arize-phoenix-otel openinference-instrumentation-langchain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install langchain_openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from getpass import getpass\n",
    "\n",
    "from langgraph.graph import END, START, StateGraph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"OPENAI_API_KEY\"] = getpass(\"🔑 Enter your OpenAI API key: \")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "64eGWAjjqJzl"
   },
   "source": [
    "# Configure Phoenix Tracing\n",
    "\n",
    "Make sure you go to https://app.phoenix.arize.com/ and generate an API key. This will allow you to trace your Langgraph application with Phoenix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if \"PHOENIX_API_KEY\" not in os.environ:\n",
    "    os.environ[\"PHOENIX_API_KEY\"] = getpass(\"🔑 Enter your Phoenix API key: \")\n",
    "\n",
    "if \"PHOENIX_COLLECTOR_ENDPOINT\" not in os.environ:\n",
    "    os.environ[\"PHOENIX_COLLECTOR_ENDPOINT\"] = getpass(\"🔑 Enter your Phoenix Collector Endpoint\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from phoenix.otel import register\n",
    "\n",
    "tracer_provider = register(project_name=\"Router\", auto_instrument=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5ttz1b_4qMo-"
   },
   "source": [
    "# LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0.3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Literal, TypedDict\n",
    "\n",
    "from langchain_core.messages import HumanMessage, SystemMessage\n",
    "from langchain_core.pydantic_v1 import BaseModel, Field"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GsxvEUl9qcac"
   },
   "source": [
    "# Router\n",
    "We define a Pydantic schema to structure the router’s output. This schema ensures the LLM returns one of three valid routing targets: \"billing\", \"tech_support\", or \"general_info\".\n",
    "We then wrap our LLM with with_structured_output, allowing LangGraph to enforce this structured response during routing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Route(BaseModel):\n",
    "    step: Literal[\"billing\", \"tech_support\", \"general_info\"] = Field(\n",
    "        description=\"Classify the support request\"\n",
    "    )\n",
    "\n",
    "\n",
    "router = llm.with_structured_output(Route)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3uW4u7YVqyUW"
   },
   "source": [
    "# Defining Graph State\n",
    "This state schema captures the lifecycle of a routed request. It stores the original user input (input), the classification decision made by the router (decision), and the final response generated by the appropriate support handler (output). Each node in the graph will read from and write to this shared state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class State(TypedDict):\n",
    "    input: str\n",
    "    decision: str\n",
    "    output: str"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XSvmCZ3qrCSj"
   },
   "source": [
    "# Support Agent Nodes: Specialized Response Handlers\n",
    "This section defines three specialized LLM-powered agents, each responsible for handling a different category of user queries:\n",
    "\n",
    "**Billing Agent**: Uses user billing history and a billing policy context to respond to invoice or refund-related questions.\n",
    "\n",
    "**Tech Support Agent**: Answers common troubleshooting queries using a predefined support knowledge base.\n",
    "\n",
    "**General Info Agent**: Responds to account, subscription, and policy-related questions using general FAQs.\n",
    "\n",
    "Each agent reads the input field from the graph's state and returns a generated output tailored to its domain. These agents form the execution endpoints of the router graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Billing Agent\n",
    "billing_history_data = [\n",
    "    {\"invoice_id\": \"INV001\", \"date\": \"2024-11-01\", \"amount\": \"$29.99\"},\n",
    "    {\"invoice_id\": \"INV002\", \"date\": \"2024-12-01\", \"amount\": \"$29.99\"},\n",
    "    {\"invoice_id\": \"INV003\", \"date\": \"2025-01-01\", \"amount\": \"$39.99\"},\n",
    "]\n",
    "\n",
    "billing_general_context = (\n",
    "    \"Billing inquiries may include refunds, invoices, plan upgrades, or charges. \"\n",
    "    \"Our system charges users monthly based on their plan. Refunds are processed within 5–7 business days.\"\n",
    ")\n",
    "\n",
    "\n",
    "def billing_agent(state: State):\n",
    "    user_input = state[\"input\"]\n",
    "    invoice_summary = \"\\n\".join(\n",
    "        f\"• Invoice {item['invoice_id']} on {item['date']} for {item['amount']}\"\n",
    "        for item in billing_history_data\n",
    "    )\n",
    "    prompt = (\n",
    "        f\"You are a helpful billing assistant.\\n\"\n",
    "        f\"Here is the user's billing history:\\n{invoice_summary}\\n\\n\"\n",
    "        f\"General billing context:\\n{billing_general_context}\\n\\n\"\n",
    "        f\"User query:\\n{user_input}\"\n",
    "    )\n",
    "    result = llm.invoke(prompt)\n",
    "    return {\"output\": result.content}\n",
    "\n",
    "\n",
    "# Tech Support Agent\n",
    "tech_support_kb = (\n",
    "    \"Common issues include login errors, app crashes, and network connectivity. \"\n",
    "    \"To fix login errors, check your email and reset your password. \"\n",
    "    \"For crashes, try reinstalling the app. If your connection is unstable, restart your router.\"\n",
    ")\n",
    "\n",
    "\n",
    "def tech_support_agent(state: State):\n",
    "    prompt = (\n",
    "        \"You are a tech support assistant. Use the knowledge base below to answer the user's question.\\n\\n\"\n",
    "        f\"Knowledge Base:\\n{tech_support_kb}\\n\\n\"\n",
    "        f\"User query:\\n{state['input']}\"\n",
    "    )\n",
    "    result = llm.invoke(prompt)\n",
    "    return {\"output\": result.content}\n",
    "\n",
    "\n",
    "# General Info Agent\n",
    "general_info_kb = (\n",
    "    \"We offer 3 subscription plans: Basic, Pro, and Enterprise. \"\n",
    "    \"Support is available 24/7. You can cancel your subscription any time from the account settings page.\"\n",
    ")\n",
    "\n",
    "\n",
    "def general_info_agent(state: State):\n",
    "    prompt = (\n",
    "        \"You are a general info assistant. Use the knowledge base below to answer the user's question.\\n\\n\"\n",
    "        f\"Knowledge Base:\\n{general_info_kb}\\n\\n\"\n",
    "        f\"User query:\\n{state['input']}\"\n",
    "    )\n",
    "    result = llm.invoke(prompt)\n",
    "    return {\"output\": result.content}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "n0mBMBNJravU"
   },
   "source": [
    "# Intent Classification & Routing Logic\n",
    "This section introduces the decision-making logic that powers the routing mechanism:\n",
    "\n",
    "**classify_intent Node**: Uses an LLM wrapped with a structured output schema to classify the user's query into one of three support categories — billing, tech_support, or general_info. The result is stored in the state's decision key.\n",
    "\n",
    "**route_to_agent Function**: A conditional router that examines the classification result and sends the request to the corresponding specialized agent node. This allows LangGraph to dynamically direct queries to the most relevant response module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def classify_intent(state: State):\n",
    "    decision = router.invoke(\n",
    "        [\n",
    "            SystemMessage(\n",
    "                content=\"Classify this support request as billing, tech_support, or general_info.\"\n",
    "            ),\n",
    "            HumanMessage(content=state[\"input\"]),\n",
    "        ]\n",
    "    )\n",
    "    return {\"decision\": decision.step}\n",
    "\n",
    "\n",
    "def route_to_agent(state: State):\n",
    "    if state[\"decision\"] == \"billing\":\n",
    "        return \"billing_agent\"\n",
    "    elif state[\"decision\"] == \"tech_support\":\n",
    "        return \"tech_support_agent\"\n",
    "    elif state[\"decision\"] == \"general_info\":\n",
    "        return \"general_info_agent\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sTJslcdprjvK"
   },
   "source": [
    "# Building the Routing Graph\n",
    "Here, we construct the complete LangGraph workflow by registering the nodes and defining their connections:\n",
    "\n",
    "**Node Registration**: All agents (billing_agent, tech_support_agent, general_info_agent) and the classify_intent node are added to the graph.\n",
    "\n",
    "**Conditional Routing**: After the graph starts at classify_intent, it uses the output of route_to_agent to forward the query to the appropriate agent node based on the classified intent.\n",
    "\n",
    "**Terminal Edges**: Each agent node directly leads to END, finalizing the workflow after responding to the query.\n",
    "\n",
    "This design reflects a typical customer support router architecture, enabling modular, extensible handling of diverse query types."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "builder = StateGraph(State)\n",
    "\n",
    "builder.add_node(\"classify_intent\", classify_intent)\n",
    "builder.add_node(\"billing_agent\", billing_agent)\n",
    "builder.add_node(\"tech_support_agent\", tech_support_agent)\n",
    "builder.add_node(\"general_info_agent\", general_info_agent)\n",
    "\n",
    "builder.add_edge(START, \"classify_intent\")\n",
    "\n",
    "builder.add_conditional_edges(\n",
    "    \"classify_intent\",\n",
    "    route_to_agent,\n",
    "    {\n",
    "        \"billing_agent\": \"billing_agent\",\n",
    "        \"tech_support_agent\": \"tech_support_agent\",\n",
    "        \"general_info_agent\": \"general_info_agent\",\n",
    "    },\n",
    ")\n",
    "\n",
    "builder.add_edge(\"billing_agent\", END)\n",
    "builder.add_edge(\"tech_support_agent\", END)\n",
    "builder.add_edge(\"general_info_agent\", END)\n",
    "\n",
    "workflow = builder.compile()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZopT2jdFrtnE"
   },
   "source": [
    "# Let's run some queries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "queries = [\n",
    "    \"Why was I charged $39.99 this month?\",  # billing\n",
    "    \"The app keeps crashing when I open it.\",  # tech support\n",
    "    \"Can I cancel my subscription anytime?\",  # general info\n",
    "    \"Can you show me all my past invoices?\",  # billing\n",
    "    \"How do I fix login issues with my account?\",  # tech support\n",
    "]\n",
    "\n",
    "# Run each query through the router\n",
    "for i, query in enumerate(queries, start=1):\n",
    "    print(f\"\\n--- Query {i} ---\")\n",
    "    state = workflow.invoke({\"input\": query})\n",
    "    print(f\"Input: {query}\")\n",
    "    print(f\"Response:\\n{state['output']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_EbruW27rxxb"
   },
   "source": [
    "# Make sure to view your traces in Phoenix!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sK7W42q2PAik"
   },
   "source": [
    "# Evals\n",
    "\n",
    "In this section we will add a simple eval that evaluates whether the router chose the correct subagent. We will use the tool calling eval template from Phoenix, as routing to sub agents mimics the process of picking a tool to call."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8EtFy2qucmgM"
   },
   "source": [
    "## Get all root spans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import phoenix as px\n",
    "from phoenix.evals import (\n",
    "    TOOL_CALLING_PROMPT_RAILS_MAP,\n",
    "    TOOL_CALLING_PROMPT_TEMPLATE,\n",
    "    OpenAIModel,\n",
    "    llm_classify,\n",
    ")\n",
    "\n",
    "df = px.Client().get_spans_dataframe(\"name == 'LangGraph'\", project_name=\"Router\")\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df.rename(columns={\"attributes.output.value\": \"tool_call\"}, inplace=True)\n",
    "df.rename(columns={\"attributes.input.value\": \"question\"}, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_clean = df[[\"context.span_id\", \"question\", \"tool_call\"]]\n",
    "df_clean"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zMGesMxCcwf7"
   },
   "source": [
    "## Add our subagent definitions as tools\n",
    "\n",
    "Each column must contain information regarding the tools that could have been chosen from. Since our \"tools\" in this case are our subagents, we pass in the definitions (descriptions) of our subagents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subagent_definitions = \"\"\"\n",
    " billing: Uses user billing history and a billing policy context to respond to invoice or refund-related questions.\n",
    "tech_support: Answers common troubleshooting queries using a predefined support knowledge base.\n",
    "general_info: Responds to account, subscription, and policy-related questions using general FAQs.\n",
    "\"\"\"\n",
    "df_clean[\"tool_definitions\"] = subagent_definitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_clean"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QDQNbzsZdBrq"
   },
   "source": [
    "# Extracting the subagent that was chosen."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "\n",
    "def extract_decision(x):\n",
    "    \"\"\"Extracts the decision value from a string using regex.\n",
    "\n",
    "    Handles cases where the regex search returns None.\n",
    "    \"\"\"\n",
    "    if x is None:\n",
    "        return None\n",
    "    match = re.search(r'\"decision\"\\s*:\\s*\"([^\"]+)\"', x)\n",
    "    if match:\n",
    "        return match.group(1)\n",
    "    else:\n",
    "        return None  # Or any default value you prefer\n",
    "\n",
    "\n",
    "df_clean[\"tool_call\"] = df_clean[\"tool_call\"].apply(extract_decision)\n",
    "df_clean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_clean.to_csv(\"router_evals.csv\", index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pARcGKpddGVF"
   },
   "source": [
    "# Running Evals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_model = OpenAIModel(model=\"gpt-4o\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rails = list(TOOL_CALLING_PROMPT_RAILS_MAP.values())\n",
    "\n",
    "response_classifications = llm_classify(\n",
    "    dataframe=df_clean,\n",
    "    template=TOOL_CALLING_PROMPT_TEMPLATE,\n",
    "    model=eval_model,\n",
    "    rails=rails,\n",
    "    provide_explanation=True,\n",
    ")\n",
    "response_classifications[\"score\"] = (response_classifications[\"label\"] == \"correct\").astype(int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response_classifications"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Sthcb_rddLOj"
   },
   "source": [
    "# Exporting Evals to Phoenix!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response_classifications.drop(\n",
    "    columns=[\"exceptions\", \"execution_status\", \"execution_seconds\"], inplace=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response_classifications.to_parquet(\"router_evals.parquet\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from phoenix.client import AsyncClient\n",
    "\n",
    "px_client = AsyncClient()\n",
    "df = response_classifications.copy()\n",
    "df.index.name = \"span_id\"\n",
    "await px_client.spans.log_span_annotations_dataframe(\n",
    "    dataframe=df,\n",
    "    annotation_name=\"Routing Eval\",\n",
    "    annotator_kind=\"LLM\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
